from openvino.inference_engine import IECore
import numpy as np
import cv2
import time

ie = IECore()
# for device in ie.available_devices:
#     print(device)

model = "best-tiny-230807.onnx"
net = ie.read_network(model=model)
input_blob = next(iter(net.input_info))
out_blob = next(iter(net.outputs))
net.batch_size = 1  # batchsize

n, c, h, w = net.input_info[input_blob].input_data.shape
print(n, c, h, w)

exec_net = ie.load_network(network=net, device_name="CPU")
image = cv2.imread("C:/Users/1234/Desktop/img_test/yolotestdata/markyangben/circle/xialiuji/1.bmp")
# 计算需要添加的上下左右边框大小
top = (1280 - 1024) // 2
bottom = 1280 - 1024 - top
left = 0
right = 0

# 在图像周围添加边框
image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(255, 255, 255))
image = image.transpose((2, 0, 1))
image = np.array(image, np.float32)
start = time.time()
res = exec_net.infer(inputs={input_blob: [image]})
output_res = res[out_blob]

class_index = np.argmax(output_res)

box_num = class_index.ndim
# print(output_res[0][0])
for obj in output_res[0][0]:
    if obj[2] > 0.7:
        print(obj[2])

if box_num > 0:
    boxs = res["boxes"]
    # for obj in boxs:
    # xmin = int(obj[0] * iw / w)
else:
    print("检测失败")

print('infer total time is %.4f s' % (time.time() - start))

#
# import cv2 as cv
# import numpy as np
# from openvino.runtime import Core
# class OpenVINOTextDetector():
#     def __init__(self):
# # 预处理设置 0 - 放缩, 1- 保持比例
# # self.preprocess_img_mode = 0
# # 插值方式, 0 - 最近邻 1 - 线性, 2 - 立方
# # self.interpolate_mode = cv.INTER_LINEAR
# # self.score_threshold = 0.5
# # self.init_text_detector()
# # def init_text_detector(self):
# # ie = Core()
# # model = ie.read_model(model="D:/python/openvm/models/text-detection-0004.xml",
# # weights="D:/python/openvm/models/text-detection-0004.bin")
# # self.compiled_model = ie.compile_model(model=model, device_name="CPU")
# # self.input_layer = next(iter(self.compiled_model.inputs))
# # model/segm_logits/add, model/link_logits_/add        # 1, 192, 320, 2
# # it = iter(self.compiled_model.outputs)
# # self.output_layer1 = next(it)
# # self.output_layer2 = next(it)
# # def format_input(self, image):        n, h, w, c = self.input_layer.shape        if self.preprocess_img_mode == 0:            resized_image = cv.resize(image, (w, h), interpolation=self.interpolate_mode)        if self.preprocess_img_mode == 1:            resized_image = np.zeros((h, w, 3), np.uint8)            rows, cols, _ = image.shape            rate_y = rows / h            rate_x = cols / w            if rate_y <= 1.0 and rate_x <= 1.0:                resized_image[0:rows, 0:cols] = image            if rate_y > 1.0 or rate_x > 1.0:                rate_max = max(rate_x, rate_y)                rh = int(rows / rate_max)                rw = int(cols / rate_max)                rimg = cv.resize(image, (rw, rh))                resized_image[0:rh, 0:rw] = rimg        input_data = np.expand_dims(resized_image, 0).astype(np.float32)        return input_data            def exec(self, image: np.ndarray) -> dict:        ih, iw, _ = image.shape        input_data = self.format_input(image)        outputs = self.compiled_model([input_data])        out1 = np.squeeze(outputs[self.output_layer1])        _, oh, ow, _ = outputs[self.output_layer1].shape        pixel_mask = np.zeros((oh, ow), dtype=np.uint8)        for row in range(oh):            for col in range(ow):                pv2 = out1[row, col, 1]                if pv2 > self.score_threshold:                    pixel_mask[row, col] = 255        mask = cv.resize(pixel_mask, (iw, ih))        contours, hiearchy = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)        text_boxes = []        for cnt in range(len(contours)):            x, y, w, h = cv.boundingRect(contours[cnt])            text_boxes.append((x, y, w, h))        return text_boxes
# if __name__ == "__main__":
#     mt = OpenVINOTextDetector()
#     image = cv.imread("D:/openvino_test.png")
#     boxes = mt.exec(image)
#     for box in boxes:
#         x, y, w, h = box
#         cv.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), 2, 8, 0)
#         cv.imshow("OpenVINO2022 Python SDK -Text Detect Demo", image)
#         cv.imwrite("D:/result.png",image)
#         cv.waitKey(0)
#         cv.destroyAllWindows()